import sys
import requests
import time
import os
import oss2

api_key = "sk-fe5d53570f2d4c2fba9cba676439f559"


def upload_image_to_oss(auth, bucket_name, endpoint, local_file_path, oss_file_path):
    """
    上传图片到阿里云 OSS
    :param auth: OSS 认证信息
    :param bucket_name: 存储桶名称
    :param endpoint: OSS 访问域名
    :param local_file_path: 本地文件路径
    :param oss_file_path: 上传到 OSS 的文件路径
    :return: 上传后的文件 URL
    """
    bucket = oss2.Bucket(auth, endpoint, bucket_name)
    result = bucket.put_object_from_file(oss_file_path, local_file_path)
    if result.status == 200:
        print(f"上传文件 {local_file_path} 到 OSS 成功")
        print(f"文件 URL: https://{bucket_name}.{endpoint}/{oss_file_path}")
        return f"https://{bucket_name}.{endpoint}/{oss_file_path}"
    else:
        print(f"上传文件 {local_file_path} 到 OSS 失败")
        return None


def submit_try_on_task(api_key, person_image_url, top_garment_url=None, bottom_garment_url=None, resolution=-1,
                       restore_face=True):
    url = "https://dashscope.aliyuncs.com/api/v1/services/aigc/image2image/image-synthesis"
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}",
        "X-DashScope-Async": "enable"
    }
    data = {
        "model": "aitryon",
        "input": {
            "person_image_url": person_image_url,
            "top_garment_url": top_garment_url,
            "bottom_garment_url": bottom_garment_url
        },
        "parameters": {
            "resolution": resolution,
            "restore_face": restore_face
        }
    }

    try:
        response = requests.post(url, headers=headers, json=data)
        response.raise_for_status()
        result = response.json()
        return result["output"]["task_id"]
    except requests.exceptions.RequestException as e:
        print(f"提交任务时发生错误: {e}")
        try:
            print("响应内容:", response.json())  # 尝试打印响应内容
        except:
            print("无法获取响应内容")
        return None


def check_task_status(api_key, task_id):
    url = f"https://dashscope.aliyuncs.com/api/v1/tasks/{task_id}"
    headers = {
        "Authorization": f"Bearer {api_key}"
    }

    try:
        response = requests.get(url, headers=headers)
        response.raise_for_status()
        result = response.json()
        print("完整的 API 响应:", result)  # 打印完整响应
        return result["output"]["task_status"], result["output"].get("image_url")
    except requests.exceptions.RequestException as e:
        print(f"查询任务状态时发生错误: {e}")
        return None, None


if __name__ == "__main__":

    person_image_url = sys.argv[1]
    top_garment_url = sys.argv[2]
    bottom_garment_url = sys.argv[3]
    # person_image_url  =  "https://shenhuixaioliweb1.oss-cn-shanghai.aliyuncs.com/test1.jpg"
    # top_garment_url = "https://shenhuixaioliweb1.oss-cn-shanghai.aliyuncs.com/trousers/00b103dcb618d22c6c89204f442ce9df.jpg"
    # bottom_garment_url = "https://shenhuixaioliweb1.oss-cn-shanghai.aliyuncs.com/1d1a97f8-32fb-4641-aecb-a23fa09b8476.jpg"

    if person_image_url and top_garment_url and bottom_garment_url:
        task_id = submit_try_on_task(api_key, person_image_url, top_garment_url, bottom_garment_url)
        if task_id:
            print(f"任务 ID: {task_id}")
            # 循环查询任务状态，直到任务完成
            while True:
                task_status, image_url = check_task_status(api_key, task_id)
                if task_status == "SUCCEEDED":
                    print(image_url)
                    break
                elif task_status == "FAILED":
                    print("任务失败")
                    break
                elif task_status is None:
                    break
                else:
                    print(f"任务状态: {task_status}，等待中...")
                    time.sleep(5)  # 每隔 5 秒查询一次
